#include "tunnel.h"

Tunnel::Tunnel(yhaida::EventLoop *loop,
               const yhaida::InetAddress &serverAddr)
            //    const yhaida::TcpConnectionPtr &serverConn)
    : _client(loop, serverAddr)
    //   _serverConn(serverConn)
{
    // LOG_INFO << "Tunnel " << serverConn->peerAddress().toHostPort()
    //          << " <-> " << serverAddr.toHostPort();
    LOG_INFO << "Tunnel "<< serverAddr.toHostPort();
}
Tunnel::~Tunnel()
{
    LOG_INFO << "~Tunnel";
}

void Tunnel::setup()
{
    _client.setConnectionCallback(std::bind(&Tunnel::onClientConnection, shared_from_this(), _1));
    _client.setMessageCallback(std::bind(&Tunnel::onClientMessage, shared_from_this(), _1, _2, _3));
    // typedef std::function<void(const TcpConnectionPtr &, size_t)> HighWaterMarkCallback;
    // void setHighWaterMarkCallback(const HighWaterMarkCallback &cb, size_t highWaterMark);
}
void Tunnel::connect()
{
    _client.connect();
}
void Tunnel::disconnect()
{
    _client.disconnect();
}

void Tunnel::teardown()
{
    _client.setConnectionCallback(yhaida::defaultConnectionCallback);
    _client.setMessageCallback(yhaida::defaultMessageCallback);
    if (_serverConn)
    {
        // _serverConn->setContext(boost::any());
        _serverConn->shutdown();
    }
    _clientConn.reset();
}
void Tunnel::setServerConn(const yhaida::TcpConnectionPtr &conn)
{
    _serverConn = conn;
    if (_serverConn)
    {
        _serverConn->setHighWaterMarkCallback(
        std::bind(&Tunnel::onHighWaterMarkWeak,
                  std::weak_ptr<Tunnel>(shared_from_this()), kServer, _1, _2),
        1024 * 1024);
        // _serverConn->setContext(conn);
        LOG_DEBUG;
        _serverConn->startRead();
        if (_serverConn->inputBuffer()->readableBytes() > 0)
        {
            conn->send(_serverConn->inputBuffer());
        }
    }
}
void Tunnel::onClientConnection(const yhaida::TcpConnectionPtr &conn)
{
    LOG_DEBUG << (conn->connected() ? "UP" : "DOWN");
    if (conn->connected())
    {
        conn->setTcpNoDelay(true);      //确保和主服务器之间减少延迟
        conn->setHighWaterMarkCallback( //控制流量
            std::bind(&Tunnel::onHighWaterMarkWeak,
                      std::weak_ptr<Tunnel>(shared_from_this()), kClient, _1, _2),
            1024 * 1024);
        // _serverConn->setContext(conn);
        // LOG_DEBUG;
        // _serverConn->startRead();
        _clientConn = conn;
        // if (_serverConn->inputBuffer()->readableBytes() > 0)
        // {
        //     conn->send(_serverConn->inputBuffer());
        // }
    }
    else
    {
        teardown();
    }
}
void Tunnel::onClientMessage(const yhaida::TcpConnectionPtr &conn,
                             yhaida::Buffer *buffer,
                             yhaida::Timestamp receiveTime)
{
    LOG_DEBUG << conn->name() << " " << buffer->readableBytes();
    if (_serverConn)
    {
        _serverConn->send(buffer);
        //FIXME: 回复消息不断开连接会复用这条TCP连接
        // _serverConn->shutdown();
    }
    else
    {
        buffer->retrieveAll();
        abort();
    }
}

void Tunnel::onHighWaterMark(ServerClient which,
                             const yhaida::TcpConnectionPtr &conn,
                             size_t bytesToSent)
{
    LOG_INFO << (which == kServer ? "server" : "client")
             << " onHighWaterMark " << conn->name()
             << " bytes " << bytesToSent;
    if (which == kServer) //代理服务器
    {
        if (_serverConn->outputBuffer()->readableBytes() > 0)
        {
            _clientConn->stopRead();
            _serverConn->setWriteCompleteCallback(
                std::bind(&Tunnel::onWriteCompleteWeak, std::weak_ptr<Tunnel>(shared_from_this()), kServer, _1));
        }
    }
    else //主服务器的客户端
    {
        if (_clientConn->outputBuffer()->readableBytes() > 0)
        {
            _serverConn->stopRead();
            _clientConn->setWriteCompleteCallback(
                std::bind(&Tunnel::onWriteCompleteWeak,
                          std::weak_ptr<Tunnel>(shared_from_this()), kClient, _1));
        }
    }
}
void Tunnel::onHighWaterMarkWeak(const std::weak_ptr<Tunnel> &wkTunnel,
                                 ServerClient which,
                                 const yhaida::TcpConnectionPtr &conn,
                                 size_t bytesToSent)
{
    std::shared_ptr<Tunnel> tunnel = wkTunnel.lock();
    if (tunnel)
    {
        tunnel->onHighWaterMark(which, conn, bytesToSent);
    }
}

void Tunnel::onWriteComplete(ServerClient which,
                             const yhaida::TcpConnectionPtr &conn)
{
    LOG_INFO << (which == kServer ? "server" : "client")
             << " onWriteComplete " << conn->name();
    if (which == kServer)
    {
        _clientConn->startRead();
        _serverConn->setWriteCompleteCallback(yhaida::WriteCompleteCallback());
    }
    else
    {
        _serverConn->startRead();
        _clientConn->setWriteCompleteCallback(yhaida::WriteCompleteCallback());
    }
}
void Tunnel::onWriteCompleteWeak(const std::weak_ptr<Tunnel> &wkTunnel,
                                 ServerClient which,
                                 const yhaida::TcpConnectionPtr &conn)
{
    std::shared_ptr<Tunnel> tunnel = wkTunnel.lock();
    if (tunnel)
    {
        tunnel->onWriteComplete(which, conn);
    }
}